use crate::Var;
use super::Optim;

pub struct Momentum {
    parameters: Vec<Var>,

    learn_rate: f64,
    momentum: f64,
    velocity: Vec<f64>,
}

impl Momentum {
    pub fn new<'a, I: Iterator<Item = &'a Var>>(parameters: I, learn_rate: f64, momentum: f64) -> Self {
        let parameters: Vec<_> = parameters.cloned().collect();
        let velocity = vec![0.; parameters.len()];
        Self { 
            parameters,
            learn_rate, 
            momentum,
            velocity
        }
    }

    pub fn default<'a, I: Iterator<Item = &'a Var>>(parameters: I) -> Self {
        Self::new(parameters, 0.01, 0.9)
    }

    pub fn learn_rate(&self) -> f64 {
        self.learn_rate
    }

    pub fn momentum(&self) -> f64 {
        self.momentum
    }
}

impl Optim for Momentum {
    fn parameters(&self) -> impl Iterator<Item = &Var> {
        self.parameters.iter()
    }

    fn step(&mut self) {
        for (param, v) in self.parameters.iter_mut().zip(self.velocity.iter_mut()) {
            // v = momentum * v + grad
            *v = self.momentum * (*v) + param.gard();
            param.increase_data(-self.learn_rate * (*v));
        }
    }
}